import random
import torch
from torch.utils.data import Dataset
from utils.graph_util import calculate_laplacian_with_self_loop


class SingleTraceDataset(Dataset):
    def __init__(self, trace_data, window_size=14, label=None):
        self.trace_data = trace_data   # (T_i, N_i)
        self.window = window_size
        self.label = label

    def __len__(self):
        return self.trace_data.shape[0] - self.window

    def __getitem__(self, idx):
        input_window = self.trace_data[idx:idx+self.window]
        output = self.trace_data[idx+self.window]  # (N_i)
        if self.label is None:
            return input_window, output, 0
        else:
            return input_window, output, self.label[idx+self.window]


class TaskDataset(Dataset):
    def __init__(self, trace_list, adj_list, window_size=14, shot=64, query=64, train_ratio=0.7):
        self.window = window_size
        self.shot = shot
        self.query = query
        self.num_type = len(trace_list)

        self.laplacian_list = []
        for adj in adj_list:
            lap = calculate_laplacian_with_self_loop(adj)
            self.laplacian_list.append(lap)

        self.train_list = []
        self.val_list = []
        self.train_choice_list = []
        self.val_choice_list = []
        self.train_len = 0
        self.val_len = 0
        for i in range(self.num_type):
            t_len = trace_list[i].shape[0]
            num_train = int(train_ratio * t_len)
            num_val = t_len - num_train
            self.train_list.append(trace_list[i][:num_train, :])
            self.val_list.append(trace_list[i][num_train:, :])
            self.train_choice_list.append([j for j in range(num_train-window_size)])
            self.val_choice_list.append([j for j in range(num_val-window_size)])
            self.train_len += (num_train - window_size)
            self.val_len += (num_val - window_size)

        self.spt_x_list = []
        self.spt_y_list = []
        self.qry_x_list = []
        self.qry_y_list = []
        for i in range(self.num_type):
            self.spt_x_list.append([])
            self.spt_y_list.append([])
            self.qry_x_list.append([])
            self.qry_y_list.append([])
            for j in range(window_size, self.train_list[i].shape[0]):
                x = self.train_list[i][j-window_size:j]   # (w, N_i)
                y = self.train_list[i][j]   # (N_i)
                self.spt_x_list[i].append(x)
                self.spt_y_list[i].append(y)
            for j in range(window_size, self.val_list[i].shape[0]):
                x = self.val_list[i][j-window_size:j]    # (w, N_i)
                y = self.val_list[i][j]   # (N_i)
                self.qry_x_list[i].append(x)
                self.qry_y_list[i].append(y)

    def sample(self, shot=None, query=None):
        if shot is None:
            shot = self.shot
        if query is None:
            query = self.query
        type_idx = random.randint(0, self.num_type-1)
        laplacian = self.laplacian_list[type_idx]                           # (N_i, N_i)

        spt_ids = random.choices(self.train_choice_list[type_idx], k=shot)
        qry_ids = random.choices(self.val_choice_list[type_idx], k=query)
        spt_x = TaskDataset._get_batch(self.spt_x_list[type_idx], spt_ids)  # (shot, T_i, N_i)
        spt_y = TaskDataset._get_batch(self.spt_y_list[type_idx], spt_ids)  # (shot, N_i)
        qry_x = TaskDataset._get_batch(self.qry_x_list[type_idx], qry_ids)  # (query, T_i, N_i)
        qry_y = TaskDataset._get_batch(self.qry_y_list[type_idx], qry_ids)  # (query, N_i)

        return (spt_x, spt_y, qry_x, qry_y, laplacian)

    @staticmethod
    def _get_batch(source_list, ids):
        tmp_list = list(map(lambda i: source_list[i], ids))
        return torch.stack(tmp_list, dim=0)